#include "./ota_service.h"
#include "../driver/cmd_uart.h"
#include "../utils/ota_info.h"
#include "../device/compat.h"

namespace service {

using namespace driver;
using namespace utils;
using namespace conn;

struct PkgReceiver {
    uint32_t pkg_index = 0;
    uint32_t pkg_size = 0;
    uint16_t pkg_id = 0;

    void init(uint32_t size) {
        pkg_index = 0;
        pkg_size = size;
        pkg_id = 0;
    }

    Error recv(uint16_t id, const ByteView &bw) {
        Error err;
        if (id != pkg_id) {
            err = ec::INVALID;
            goto finish;
        }
        pkg_id ++;
        
        err = flash::write(flash::OTA_DATA_ADDR + pkg_index, bw.buf, bw.len);
        if (err) {
            goto finish;
        }

        pkg_index += bw.len;
    finish:
        return err;
    }

    bool is_complete() {
        return pkg_index == pkg_size;
    }
};

static PkgReceiver pkg_receiver;

void ota_start(const Message &msg) {
    uint32_t pkg_size;
    Error err = msg.parse_req(pkg_size);
    if (err) {
        goto finish;
    }
    if (pkg_size > flash::OTA_DATA_SIZE) {
        err = ec::INVALID;
        goto finish;
    }
    err = flash::erase(flash::OTA_DATA_ADDR, pkg_size);
    pkg_receiver.init(pkg_size);
finish:
    cmdUart->write_resp(msg.seq(), msg.req_cmd(), err);
}

void ota_translate(const Message &msg) {
    uint16_t pkg_id;
    ByteView bw;
    Error err = msg.parse_req(pkg_id, bw);
    if (err) {
        goto finish;
    }
    err = pkg_receiver.recv(pkg_id, bw);
    if (err) {
        goto finish;
    }
finish:
    cmdUart->write_resp(msg.seq(), msg.req_cmd(), err);
}

void ota_complete(const Message &msg) {
    ByteView bw;
    uint8_t md5_hash[16];
    Error err = msg.parse_req(bw);
    if (err) {
        goto finish;
    }
    if ((bw.len != 16) || (!pkg_receiver.is_complete())) {
        err = ec::INVALID;
        goto finish;
    }
    md5::sum(reinterpret_cast<const void *>(flash::OTA_DATA_ADDR), pkg_receiver.pkg_size, md5_hash);
    if (memcmp(md5_hash, bw.buf, bw.len) != 0) {
        err = ec::INVALID;
        goto finish;
    }
    ota_info->init(pkg_receiver.pkg_size);
    err = ota_info.save();
finish:
    cmdUart->write_resp(msg.seq(), msg.req_cmd(), err);
    osDelay(500);
    NVIC_SystemReset();
}

}
